#!/usr/bin/env python3

# Copyright (c) 2021-2024, Camptocamp SA
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# The views and conclusions contained in the software and documentation are those
# of the authors and should not be interpreted as representing official policies,
# either expressed or implied, of the FreeBSD Project.

import argparse
import shlex
import subprocess
import sys


def main() -> None:
    """Restore the database backup."""
    parser = argparse.ArgumentParser(description="Restore the database backup.")
    parser.add_argument("--verbose", action="store_true", help="Print the command that is executed.")
    parser.add_argument(
        "--dry-run", action="store_true", help="Only print the command that would be executed."
    )
    parser.add_argument(
        "--env",
        help="The env file to use to get the connection settings, by default we log in the composition to "
        "do the restore.",
    )
    parser.add_argument(
        "--docker-compose-file",
        action="append",
        default=[],
        help="The docker compose file to used.",
    )
    parser.add_argument(
        "--arg",
        action="append",
        default=[],
        help="Additional pg_restore argument, example: '--arg=--schema=main', may be used multiple times.",
    )
    parser.add_argument("file", help="The file path to restore")
    args = parser.parse_args()

    with open(args.file, "rb") as file_in:
        if args.env:
            env = {}
            with open(args.env, encoding="utf-8") as file_env:
                for line in file_env:
                    if line and line[0] != "#":
                        try:
                            index = line.index("=")
                            env[line[:index].strip()] = line[index + 1 :].strip()
                        except ValueError:
                            # Ignore lines that don't have a '='
                            pass
            command1 = [
                "docker",
                "run",
                "--rm",
                "--interactive",
                "--env=PGHOST={}".format(env["PGHOST"]),
                "--env=PGPORT={}".format(env["PGPORT"]),
                "--env=PGUSER={}".format(env["PGUSER"]),
            ]
            command2 = [
                "--env=PGPASSWORD={}".format(env["PGPASSWORD"]),
            ]
            command2_annon = ["--env=PGPASSWORD=***"]
            command3 = [
                *(["--env=PGSSLMODE={}".format(env["PGSSLMODE"])] if "PGSSLMODE" in env else []),
                *(["--env=PGOPTIONS={}".format(env["PGOPTIONS"])] if "PGOPTIONS" in env else []),
                "camptocamp/postgres:{}".format(env["POSTGRES_TAG"]),
                "pg_restore",
                "--dbname={}".format(env["PGDATABASE"]),
                *args.arg,
            ]
            if args.verbose or args.dry_run:
                subprocess.run(
                    [
                        "docker",
                        "run",
                        "--rm",
                        "camptocamp/postgres:{}".format(env["POSTGRES_TAG"]),
                        "pg_restore",
                        "--version",
                    ]
                )
                print(shlex.join([*command1, *command2_annon, *command3]))
            if args.dry_run:
                sys.exit(0)
            sys.exit(subprocess.run([*command1, *command2, *command3], stdin=file_in).returncode)
        else:
            command = [
                "docker",
                "compose",
                *[f"--file={f}" for f in args.docker_compose_file],
                "exec",
                "-T",
                "tools",
                "bash",
                "-c",
                'pg_restore --dbname="$PGDATABASE" ' + " ".join(args.arg),
            ]
            if args.verbose or args.dry_run:
                subprocess.run(["docker", "compose", "exec", "tools", "pg_restore", "--version"])
                print(shlex.join(command))
            if args.dry_run:
                sys.exit(0)
            sys.exit(subprocess.run(command, stdin=file_in).returncode)


if __name__ == "__main__":
    main()
